package com.github.netty.util;

import javax.servlet.annotation.WebServlet;
import java.io.File;
import java.io.IOException;
import java.net.JarURLConnection;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class ServletAnnotationScanner2 {

    public Set<Class<?>> scanPackage(String packageName) throws IOException, ClassNotFoundException {
        Set<Class<?>> classes = new HashSet<>();
        String path = packageName.replace('.', '/');
        ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
        Enumeration<URL> resources = classLoader.getResources(path);
        while (resources.hasMoreElements()) {
            URL resource = resources.nextElement();
            if (resource.getProtocol().equals("jar")) {
                JarURLConnection connection = (JarURLConnection) resource.openConnection();
                JarFile jarFile = connection.getJarFile();
                Enumeration<JarEntry> entries = jarFile.entries();
                while (entries.hasMoreElements()) {
                    JarEntry entry = entries.nextElement();
                    if (entry.isDirectory()) {
                        continue;
                    }
                    String className = entry.getName().replace('/', '.').substring(0, entry.getName().length() - 6);
                    if (className.startsWith(packageName)) {
                        try {
                            Class<?> clazz = classLoader.loadClass(className);
                            if (clazz.isAnnotationPresent(WebServlet.class)) {
                                classes.add(clazz);
                            }
                        } catch (NoClassDefFoundError | ClassNotFoundException e) {
                            // Log and handle the error appropriately
                        }
                    }
                }
            } else {
                URI uri = null;
                try {
                    uri = resource.toURI();
                } catch (URISyntaxException e) {
                    throw new IOException(e);
                }
                File folder = new File(uri);
                File[] files = folder.listFiles();
                if (files != null) {
                    for (File file : files) {
                        if (file.isDirectory()) {
                            continue;
                        }
                        String className = packageName + '.' + file.getName().substring(0, file.getName().length() - 6);
                        try {
                            Class<?> clazz = classLoader.loadClass(className);
                            if (clazz.isAnnotationPresent(WebServlet.class)) {
                                classes.add(clazz);
                            }
                        } catch (NoClassDefFoundError | ClassNotFoundException e) {
                            // Log and handle the error appropriately
                        }
                    }
                }
            }
        }
        return classes;
    }

    public static void main(String[] args) {
        ServletAnnotationScanner2 scanner = new ServletAnnotationScanner2();
        try {
            Set<Class<?>> servletClasses = scanner.scanPackage("your.package.name");
            for (Class<?> clazz : servletClasses) {
                System.out.println("Found @WebServlet class: " + clazz.getName());
            }
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
    }
}
